{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "# 基于Fourier Neural Operator的Burgers' equation求解\n",
                "\n",
                "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/data_driven/mindspore_burgers_FNO1D.ipynb)&emsp;[![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/zh_cn/data_driven/mindspore_burgers_FNO1D.py)&emsp;[![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/data_driven/burgers_FNO1D.ipynb)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 概述\n",
                "\n",
                "计算流体力学是21世纪流体力学领域的重要技术之一，其通过使用数值方法在计算机中对流体力学的控制方程进行求解，从而实现流动的分析、预测和控制。传统的有限元法（finite element method，FEM）和有限差分法（finite difference method，FDM）常用于复杂的仿真流程（物理建模、网格划分、数值离散、迭代求解等）和较高的计算成本，往往效率低下。因此，借助AI提升流体仿真效率是十分必要的。\n",
                "\n",
                "近年来，随着神经网络的迅猛发展，为科学计算提供了新的范式。经典的神经网络是在有限维度的空间进行映射，只能学习与特定离散化相关的解。与经典神经网络不同，傅里叶神经算子（Fourier Neural Operator，FNO）是一种能够学习无限维函数空间映射的新型深度学习架构。该架构可直接学习从任意函数参数到解的映射，用于解决一类偏微分方程的求解问题，具有更强的泛化能力。更多信息可参考[Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)。\n",
                "\n",
                "本案例教程介绍利用傅里叶神经算子的1-d Burgers方程求解方法。"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 伯格斯方程（Burgers' equation）\n",
                "\n",
                "一维伯格斯方程（1-d Burgers' equation）是一个非线性偏微分方程，具有广泛应用，包括一维粘性流体流动建模。它的形式如下：\n",
                "\n",
                "$$\n",
                "\\partial_t u(x, t)+\\partial_x (u^2(x, t)/2)=\\nu \\partial_{xx} u(x, t), \\quad x \\in(0,1), t \\in(0, 1]\n",
                "$$\n",
                "\n",
                "$$\n",
                "u(x, 0)=u_0(x), \\quad x \\in(0,1)\n",
                "$$\n",
                "\n",
                "其中$u$表示速度场，$u_0$表示初始条件，$\\nu$表示粘度系数。\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 问题描述\n",
                "\n",
                "本案例利用Fourier Neural Operator学习初始状态到下一时刻状态的映射，实现一维Burgers'方程的求解：\n",
                "\n",
                "$$\n",
                "u_0 \\mapsto u(\\cdot, 1)\n",
                "$$"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 技术路径\n",
                "\n",
                "MindFlow求解该问题的具体流程如下：\n",
                "\n",
                "1. 创建数据集。\n",
                "2. 构建模型。\n",
                "3. 优化器与损失函数。\n",
                "4. 模型训练。"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Fourier Neural Operator\n",
                "\n",
                "Fourier Neural Operator模型构架如下图所示。图中$w_0(x)$表示初始涡度，通过Lifting Layer实现输入向量的高维映射，然后将映射结果作为Fourier Layer的输入，进行频域信息的非线性变换，最后由Decoding Layer将变换结果映射至最终的预测结果$w_1(x)$。\n",
                "\n",
                "Lifting Layer、Fourier Layer以及Decoding Layer共同组成了Fourier Neural Operator。\n",
                "\n",
                "![Fourier Neural Operator模型构架](images/FNO.png)\n",
                "\n",
                "Fourier Layer网络结构如下图所示。图中V表示输入向量，上框表示向量经过傅里叶变换后，经过线性变换R，过滤高频信息，然后进行傅里叶逆变换；另一分支经过线性变换W，最后通过激活函数，得到Fourier Layer输出向量。\n",
                "\n",
                "![Fourier Layer网络结构](images/FNO-2.png)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "import os\n",
                "import time\n",
                "import numpy as np\n",
                "\n",
                "from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
                "from mindspore import context, nn, Tensor, set_seed, ops, data_sink, jit, save_checkpoint\n",
                "from mindspore import dtype as mstype\n",
                "from mindflow import FNO1D, RelativeRMSELoss, load_yaml_config, get_warmup_cosine_annealing_lr\n",
                "from mindflow.pde import UnsteadyFlowWithLoss\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "下述`src`包可以在[applications/data_driven/burgers/fno1d/src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/data_driven/burgers/fno1d/src)下载。"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "from src import create_training_dataset\n",
                "\n",
                "set_seed(0)\n",
                "np.random.seed(0)\n",
                "\n",
                "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\", device_id=0)\n",
                "use_ascend = context.get_context(attr_key='device_target') == \"Ascend\""
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "从`config`中获得模型、数据、优化器的参数。"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "config = load_yaml_config('./configs/fno1d.yaml')\n",
                "data_params = config[\"data\"]\n",
                "model_params = config[\"model\"]\n",
                "optimizer_params = config[\"optimizer\"]"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 创建数据集\n",
                "\n",
                "下载训练与测试数据集: [data_driven/burgers/fno1d/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/burgers/dataset/)。\n",
                "\n",
                "本案例根据Zongyi Li在 [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下：\n",
                "基于周期性边界，生成满足如下分布的初始条件$u_0(x)$：\n",
                "\n",
                "$$\n",
                "u_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,625(-\\Delta+25 I)^{-2}\\right)\n",
                "$$\n",
                "\n",
                "本案例选取粘度系数$\\nu=0.1$，并使用分步法求解方程，其中热方程部分在傅里叶空间中精确求解，然后使用前向欧拉方法求解非线性部分。训练集样本量为1000个，测试集样本量为200个。\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Data preparation finished\n",
                        "input_path:  (1000, 1024, 1)\n",
                        "label_path:  (1000, 1024)\n"
                    ]
                }
            ],
            "source": [
                "# create training dataset\n",
                "train_dataset = create_training_dataset(data_params, model_params, shuffle=True)\n",
                "\n",
                "# create test dataset\n",
                "test_input, test_label = np.load(os.path.join(data_params[\"root_dir\"], \"test/inputs.npy\")), \\\n",
                "                         np.load(os.path.join(data_params[\"root_dir\"], \"test/label.npy\"))\n",
                "test_input = Tensor(np.expand_dims(test_input, -2), mstype.float32)\n",
                "test_label = Tensor(np.expand_dims(test_label, -2), mstype.float32)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 构建模型\n",
                "\n",
                "网络由1层Lifting layer、1层Decoding layer以及多层Fourier Layer叠加组成：\n",
                "\n",
                "- Lifting layer对应样例代码中`FNO1D.fc0`，将输出数据$x$映射至高维；\n",
                "\n",
                "- 多层Fourier Layer的叠加对应样例代码中`FNO1D.fno_seq`，本案例采用离散傅里叶变换实现时域与频域的转换；\n",
                "\n",
                "- Decoding layer对应代码中`FNO1D.fc1`与`FNO1D.fc2`，获得最终的预测值。\n",
                "\n",
                "基于上述网络结构，进行模型初始化，其中模型参数可在[配置文件](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)中修改。"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n"
                    ]
                }
            ],
            "source": [
                "model = FNO1D(in_channels=model_params[\"in_channels\"],\n",
                "              out_channels=model_params[\"out_channels\"],\n",
                "              n_modes=model_params[\"modes\"],\n",
                "              resolutions=model_params[\"resolutions\"],\n",
                "              hidden_channels=model_params[\"hidden_channels\"],\n",
                "              n_layers=model_params[\"depths\"],\n",
                "              projection_channels=4*model_params[\"hidden_channels\"],\n",
                "              )\n",
                "model_params_list = []\n",
                "for k, v in model_params.items():\n",
                "    model_params_list.append(f\"{k}:{v}\")\n",
                "model_name = \"_\".join(model_params_list)\n",
                "print(model_name)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 优化器与损失函数\n",
                "\n",
                "使用相对均方根误差作为网络训练损失函数："
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "steps_per_epoch = train_dataset.get_dataset_size()\n",
                "lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params[\"learning_rate\"],\n",
                "                                    last_epoch=optimizer_params[\"epochs\"],\n",
                "                                    steps_per_epoch=steps_per_epoch,\n",
                "                                    warmup_epochs=1)\n",
                "optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n",
                "\n",
                "if use_ascend:\n",
                "    from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
                "    loss_scaler = DynamicLossScaler(1024, 2, 100)\n",
                "    auto_mixed_precision(model, 'O1')\n",
                "else:\n",
                "    loss_scaler = None"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## 模型训练\n",
                "\n",
                "使用 **MindSpore version >= 2.0.0**, 我们可以使用函数式编程来训练神经网络。 `MindFlow` 为非稳态问题 `UnsteadyFlowWithLoss` 提供了一个训练接口，用于模型训练和评估."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "pid: 1993\n",
                        "2023-02-01 12:14:12.2323\n",
                        "use_ascend: False\n",
                        "device_id: 4\n",
                        "Data preparation finished\n",
                        "input_path:  (1000, 1024, 1)\n",
                        "label_path:  (1000, 1024)\n",
                        "name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n",
                        "./summary/name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n",
                        "epoch: 1 train loss: 2.167046070098877 epoch time: 21.75s\n",
                        "epoch: 2 train loss: 0.5935954451560974 epoch time: 5.53s\n",
                        "epoch: 3 train loss: 0.7349425554275513 epoch time: 5.46s\n",
                        "epoch: 4 train loss: 0.6338694095611572 epoch time: 4.95s\n",
                        "epoch: 5 train loss: 0.3174982964992523 epoch time: 5.57s\n",
                        "epoch: 6 train loss: 0.3099440038204193 epoch time: 5.71s\n",
                        "epoch: 7 train loss: 0.3117891848087311 epoch time: 5.22s\n",
                        "epoch: 8 train loss: 0.1810857653617859 epoch time: 4.82s\n",
                        "epoch: 9 train loss: 0.1386510729789734 epoch time: 4.97s\n",
                        "epoch: 10 train loss: 0.2102256715297699 epoch time: 4.85s\n",
                        "================================Start Evaluation================================\n",
                        "mean rms_error: 0.027940063\n",
                        "=================================End Evaluation=================================\n",
                        "...\n",
                        "epoch: 91 train loss: 0.019643772393465042 epoch time: 4.40s\n",
                        "epoch: 92 train loss: 0.0641067773103714 epoch time: 5.48s\n",
                        "epoch: 93 train loss: 0.02199840545654297 epoch time: 5.55s\n",
                        "epoch: 94 train loss: 0.024467874318361282 epoch time: 6.24s\n",
                        "epoch: 95 train loss: 0.025712188333272934 epoch time: 5.43s\n",
                        "epoch: 96 train loss: 0.02247200347483158 epoch time: 6.48s\n",
                        "epoch: 97 train loss: 0.026637140661478043 epoch time: 6.30s\n",
                        "epoch: 98 train loss: 0.030040305107831955 epoch time: 5.16s\n",
                        "epoch: 99 train loss: 0.02589748054742813 epoch time: 5.36s\n",
                        "epoch: 100 train loss: 0.028599221259355545 epoch time: 5.90s\n",
                        "================================Start Evaluation================================\n",
                        "mean rms_error: 0.0037017763\n",
                        "=================================End Evaluation=================================\n"
                    ]
                }
            ],
            "source": [
                "problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format=\"NHWTC\")\n",
                "\n",
                "summary_dir = os.path.join(config[\"summary\"][\"summary_dir\"], model_name)\n",
                "print(summary_dir)\n",
                "\n",
                "def forward_fn(data, label):\n",
                "    loss = problem.get_loss(data, label)\n",
                "    if use_ascend:\n",
                "        loss = loss_scaler.scale(loss)\n",
                "    return loss\n",
                "\n",
                "grad_fn = ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n",
                "\n",
                "@jit\n",
                "def train_step(data, label):\n",
                "    loss, grads = grad_fn(data, label)\n",
                "    if use_ascend:\n",
                "        loss = loss_scaler.unscale(loss)\n",
                "        if all_finite(grads):\n",
                "            grads = loss_scaler.unscale(grads)\n",
                "    loss = ops.depend(loss, optimizer(grads))\n",
                "    return loss\n",
                "\n",
                "sink_process = data_sink(train_step, train_dataset, 1)\n",
                "summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n",
                "ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n",
                "if not os.path.exists(ckpt_dir):\n",
                "    os.makedirs(ckpt_dir)\n",
                "\n",
                "for epoch in range(1, optimizer_params[\"epochs\"] + 1):\n",
                "    model.set_train()\n",
                "    local_time_beg = time.time()\n",
                "    for _ in range(steps_per_epoch):\n",
                "        cur_loss = sink_process()\n",
                "    print(\n",
                "        f\"epoch: {epoch} train loss: {cur_loss.asnumpy():.8f}\"\\\n",
                "        f\" epoch time: {time.time() - local_time_beg:.2f}s\"\\\n",
                "        f\" step time: {(time.time() - local_time_beg)/steps_per_epoch:.4f}s\")\n",
                "\n",
                "    if epoch % config['summary']['test_interval'] == 0:\n",
                "        model.set_train(False)\n",
                "        print(\"================================Start Evaluation================================\")\n",
                "        rms_error = problem.get_loss(test_input, test_label)/test_input.shape[0]\n",
                "        print(f\"mean rms_error: {rms_error}\")\n",
                "        print(\"=================================End Evaluation=================================\")\n",
                "        save_checkpoint(model, os.path.join(ckpt_dir, model_params[\"name\"] + '_epoch' + str(epoch)))"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.9.19"
        },
        "vscode": {
            "interpreter": {
                "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 1
}
